#pragma once

#include <string>

#include "DeviceShare.h"

#include <infiniband/verbs.h>
#include <infiniband/mlx5_api.h>
#include <libflexio/flexio.h>

#include <bits/stdc++.h>
#include <random>
#include <sstream>

#include <sys/socket.h>
#include <arpa/inet.h>

#include <doca_buf.h>
#include <doca_buf_inventory.h>
#include <doca_ctx.h>
#include <doca_dev.h>
#include <doca_dma.h>
#include <doca_error.h>
#include <doca_log.h>
#include <doca_mmap.h>
#include <doca_argp.h>
#include <doca_sync_event.h>
#include <doca_types.h>
#include <doca_dpdk.h>
#include <doca_compress.h>
#include <doca_erasure_coding.h>
#include <doca_pcc.h>
#include <doca_pe.h>
#include <doca_rdma.h>
#include <doca_rdma_bridge.h>
#include <execinfo.h>
#include <glog/logging.h>

namespace FLEX {
    template <typename A>
    class Base {
    public:
        virtual A *get_inner_ptr() = 0;

        virtual std::string get_type_name() const = 0;

        virtual ~Base() noexcept(false) = default;
    };

    class Context: public Base<::ibv_context> {
    public:
        explicit Context(std::string device_name);

        ~Context() noexcept(false) override;

        ::ibv_context *get_inner_ptr() override;

        std::string get_type_name() const override;

        ::flexio_uar *get_flexio_uar();

        uint32_t get_flexio_uar_id();

        ::ibv_pd *get_pd();

        ::flexio_process *get_process();

        ::flexio_window *get_window();

        uint32_t get_window_id();

        void alloc_pd();

        void *alloc_memalign(size_t size, size_t align_size);

        void create_process(flexio_app *app_device);

        void generate_flexio_uar();

        void create_window();

        ::ibv_mr *get_mr(size_t index);

        ibv_mr *reg_mr(void *addr, size_t size, int flag = IBV_ACCESS_LOCAL_WRITE);

        void print_init(const std::string &f_name_prefix, int num_threads);

        void flush();

        template <typename T>
        flexio_uintptr_t move_to_dev(const T &arg) {
            flexio_uintptr_t daddr;
            flexio_status ret = flexio_copy_from_host(process, const_cast<T *>(&arg), sizeof(arg), &daddr);
            assert(ret == FLEXIO_STATUS_SUCCESS);
            return daddr;
        }

    private:
        ::ibv_context *inner_ibv_context{};
        ::ibv_pd *pd{};
        // ::flexio_prm_hca_caps *hca_caps{};
        ::flexio_process *process{};
        ::flexio_uar *flexio_uar{};
        // ::flexio_outbox *outbox{};
        std::vector<::flexio_event_handler *> event_handler_vec{};
        ::flexio_window *window{};
        std::vector<::ibv_mr *> mr_vec{};
        ::flexio_msg_stream *default_stream{};
        ::flexio_msg_stream **streams{};
        int number_of_threads;
    };

    class CommandQ: public Base<::flexio_cmdq> {
    public:
        explicit CommandQ(flexio_process *process, int num_threads, int batch_size);

        ~CommandQ() noexcept(false) override;

        ::flexio_cmdq *get_inner_ptr() override;

        std::string get_type_name() const override;

        void add_task(flexio_func_t func, flexio_uintptr_t ddata);

        void run();

        long wait_run(int time_out_seconds);

        ::flexio_cmdq *cmd_q{};
    };
}